/*
ref:
    https://blog.csdn.net/chenyefei/article/details/52314290
    https://www.jianshu.com/p/ca83254acec4
*/

#include <stdio.h>
#include <stdlib.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>

int save_yuv(AVFrame *pfrm, AVFrame *pyuv_frm, AVPacket *ppkt, AVCodecContext *pcodec_ctx, FILE *fp, struct SwsContext *psws_ctx, int *p_cnt, int *total_cnt) {
    int ret = 0;
    if ((ret = avcodec_send_packet(pcodec_ctx, ppkt)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_packet ERROR %s \n", av_err2str(ret));
        return -1;
    } else {
        while (ret >= 0) {
            ret = avcodec_receive_frame(pcodec_ctx, pfrm);
            if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN)) {
                break;
            }
            else if (ret < 0) {
                av_log(NULL, AV_LOG_ERROR, "avcodec_receive_frame ERROR %s \n", av_err2str(ret));
                return -2;
            }
                    
            sws_scale(psws_ctx, (const uint8_t* const *)pfrm->data, pfrm->linesize, 0, pfrm->height, pyuv_frm->data, pyuv_frm->linesize);
            fwrite(pyuv_frm->data[0], 1, pcodec_ctx->width * pcodec_ctx->height, fp);
            fwrite(pyuv_frm->data[1], 1, pcodec_ctx->width * pcodec_ctx->height / 4, fp);
            fwrite(pyuv_frm->data[2], 1, pcodec_ctx->width * pcodec_ctx->height / 4, fp);

            (*p_cnt)++ ;
            fprintf(stdout, "Complete frame: %d ,coded_n: %d , total: %d \n", *p_cnt, pfrm->coded_picture_number, *total_cnt);
                    
            av_frame_unref(pfrm);
        }  
    }
    return 0;
}

int main (int argc, char **argv) {
    
    if (argc != 3) {
        av_log(NULL, AV_LOG_ERROR, "You must input the input mp4 file and output yuv file! \n");
        return -1;
    }

    char *mp4name = NULL;
    char *yuvname = NULL;
    int ret = 0;
    FILE *yuv_fp = NULL;
    AVFormatContext *pfmt_ctx = NULL;
    AVStream *pstm = NULL;
    AVPacket *ppkt = NULL;
    AVFrame *pfrm = NULL;
    AVFrame *pyuv_frm = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    AVCodec *pcodec = NULL;
    struct SwsContext *psws_ctx = NULL;
    int video_idx = -1;
    int audio_idx = -1;
    int count = 0;
    int total_count = 0;

    mp4name = argv[1];
    yuvname = argv[2];

    pfmt_ctx = avformat_alloc_context();
    if (pfmt_ctx == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR  \n");
        return AVERROR(ENOMEM);
    }

    if ((ret = avformat_open_input(&pfmt_ctx, mp4name, NULL, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_open_input ERROR %s \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    if ((ret = avformat_find_stream_info(pfmt_ctx, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_find_stream_info ERROR %s \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    av_dump_format(pfmt_ctx, 0, mp4name, 0);

    for (int i=0; i<pfmt_ctx->nb_streams; i++) {
        if (pfmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_idx = i;
        } else if (pfmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            audio_idx = i;
        }
    }
    
    if (video_idx == -1 || audio_idx == -1) {
        fprintf(stderr, "Can not find right video or audio index \n");
        return -1;
    }

    pstm = pfmt_ctx->streams[video_idx];

    if ((pcodec = avcodec_find_decoder(pstm->codecpar->codec_id)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_decoder ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((pcodec_ctx = avcodec_alloc_context3(pcodec)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((ret = avcodec_parameters_to_context(pcodec_ctx, pstm->codecpar)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_from_context ERROR %s \n", av_err2str(ret));
        return AVERROR(ret);
    }

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR %s \n", av_err2str(ret));
        return AVERROR(ret);
    }

    if ((ppkt = av_packet_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((pfrm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((pyuv_frm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((yuv_fp = fopen(yuvname, "wb")) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "fopen ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((psws_ctx = sws_getContext(pcodec_ctx->width, pcodec_ctx->height, pcodec_ctx->pix_fmt, pcodec_ctx->width, pcodec_ctx->height, AV_PIX_FMT_YUV420P, SWS_FAST_BILINEAR, NULL, NULL, NULL)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "sws_getContext ERROR \n");
        return AVERROR(EINVAL);
    }
    
    av_image_alloc(pyuv_frm->data, pyuv_frm->linesize, pcodec_ctx->width, pcodec_ctx->height, AV_PIX_FMT_YUV420P, 1);

    total_count = (int) pstm->nb_frames;

    while ((ret = av_read_frame(pfmt_ctx, ppkt)) >= 0) {
        if (ppkt->stream_index == video_idx) {
            ret = save_yuv(pfrm, pyuv_frm, ppkt, pcodec_ctx, yuv_fp, psws_ctx, &count, &total_count);
        }
        av_packet_unref(ppkt);
        if (ret < 0) {
            break;
        }
    }

    save_yuv(pfrm, pyuv_frm, NULL, pcodec_ctx, yuv_fp, psws_ctx, &count, &total_count);

fail:
    fclose(yuv_fp);
    av_frame_free(&pfrm);
    av_packet_free(&ppkt);
    sws_freeContext(psws_ctx);
    avcodec_free_context(&pcodec_ctx);
    avformat_close_input(&pfmt_ctx);
    avformat_free_context(pfmt_ctx);

    return 0;
}
